#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import rospy
import numpy as np
from sensor_msgs.msg import PointCloud2, PointField
from std_msgs.msg import Header
import struct
import math

class GridGroundFilter:
    def __init__(self):
        rospy.init_node('grid_ground_filter', anonymous=True)
        
        # 参数设置
        self.grid_size = rospy.get_param('~grid_size', 0.5)  # 网格大小，默认0.5m
        self.height_threshold = rospy.get_param('~height_threshold', 0.2)  # 高度极差阈值，默认0.1m
        self.ground_intensity = rospy.get_param('~ground_intensity', 10.0)  # 地面点强度
        self.non_ground_intensity = rospy.get_param('~non_ground_intensity', 30.0)  # 非地面点强度
        
        # 订阅和发布
        self.subscriber = rospy.Subscriber('/ground_segmentation/no_ground', PointCloud2, self.pointcloud_callback)
        self.ground_pub = rospy.Publisher('/ground_points', PointCloud2, queue_size=1)
        self.non_ground_pub = rospy.Publisher('/non_ground_points', PointCloud2, queue_size=1)
        
        rospy.loginfo("Grid Ground Filter initialized with:")
        rospy.loginfo("  Grid size: %.2f m", self.grid_size)
        rospy.loginfo("  Height threshold: %.2f m", self.height_threshold)
        rospy.loginfo("  Ground intensity: %.1f", self.ground_intensity)
        rospy.loginfo("  Non-ground intensity: %.1f", self.non_ground_intensity)
        
    def pointcloud_callback(self, msg):
        """处理点云消息"""
        try:
            # 解析点云数据
            points = self.parse_pointcloud(msg)
            if len(points) == 0:
                return
                
            # 进行网格化地面分割
            ground_points, non_ground_points = self.grid_ground_segmentation(points)
            
            # 发布结果
            self.publish_pointcloud(ground_points, self.ground_pub, msg.header, self.ground_intensity)
            self.publish_pointcloud(non_ground_points, self.non_ground_pub, msg.header, self.non_ground_intensity)
            
            rospy.loginfo("Processed %d points: %d ground, %d non-ground", 
                         len(points), len(ground_points), len(non_ground_points))
                         
        except Exception as e:
            rospy.logerr("Error processing pointcloud: %s", str(e))
    
    def parse_pointcloud(self, msg):
        """解析PointCloud2消息为numpy数组"""
        # 获取点云字段信息
        fields = {field.name: field for field in msg.fields}
        
        # 检查必要的字段
        if 'x' not in fields or 'y' not in fields or 'z' not in fields:
            rospy.logwarn("PointCloud2 missing required fields (x, y, z)")
            return np.array([])
        
        # 计算点云数据
        point_step = msg.point_step
        row_step = msg.row_step
        
        # 解析点云数据
        points = []
        for i in range(msg.width * msg.height):
            offset = i * point_step
            
            # 读取x, y, z坐标
            x = struct.unpack('f', msg.data[offset + fields['x'].offset:offset + fields['x'].offset + 4])[0]
            y = struct.unpack('f', msg.data[offset + fields['y'].offset:offset + fields['y'].offset + 4])[0]
            z = struct.unpack('f', msg.data[offset + fields['z'].offset:offset + fields['z'].offset + 4])[0]
            
            # 过滤无效点
            if not (math.isnan(x) or math.isnan(y) or math.isnan(z) or 
                   math.isinf(x) or math.isinf(y) or math.isinf(z)):
                points.append([x, y, z])
        
        return np.array(points)
    
    def grid_ground_segmentation(self, points):
        """基于2D网格的地面分割"""
        if len(points) == 0:
            return [], []
        
        # 计算网格边界
        x_min, x_max = np.min(points[:, 0]), np.max(points[:, 0])
        y_min, y_max = np.min(points[:, 1]), np.max(points[:, 1])
        
        # 计算网格数量
        grid_x_count = int((x_max - x_min) / self.grid_size) + 1
        grid_y_count = int((y_max - y_min) / self.grid_size) + 1
        
        # 创建网格字典，存储每个网格内的点
        grids = {}
        
        # 将点分配到网格
        for point in points:
            x, y, z = point
            grid_x = int((x - x_min) / self.grid_size)
            grid_y = int((y - y_min) / self.grid_size)
            grid_key = (grid_x, grid_y)
            
            if grid_key not in grids:
                grids[grid_key] = []
            grids[grid_key].append(point)
        
        # 分类地面点和非地面点
        ground_points = []
        non_ground_points = []
        
        for grid_key, grid_points in grids.items():
            if len(grid_points) == 0:
                continue
                
            grid_points = np.array(grid_points)
            heights = grid_points[:, 2]
            
            # 计算高度极差
            height_range = np.max(heights) - np.min(heights)
            min_height = np.min(heights)
            
            if height_range > self.height_threshold:
                # 极差大于阈值，将最低点高度+0.1以下的点分类为地面点
                ground_height_threshold = min_height + self.height_threshold
                
                for point in grid_points:
                    if point[2] <= ground_height_threshold:
                        ground_points.append(point)
                    else:
                        non_ground_points.append(point)
            else:
                # 极差小于阈值，全部分类为地面点
                ground_points.extend(grid_points)
        
        return np.array(ground_points), np.array(non_ground_points)
    
    def publish_pointcloud(self, points, publisher, header, intensity):
        """发布点云消息"""
        if len(points) == 0:
            return
            
        # 创建PointCloud2消息
        msg = PointCloud2()
        msg.header = header
        msg.header.frame_id = header.frame_id
        
        # 设置点云字段
        msg.fields = [
            PointField('x', 0, PointField.FLOAT32, 1),
            PointField('y', 4, PointField.FLOAT32, 1),
            PointField('z', 8, PointField.FLOAT32, 1),
            PointField('intensity', 12, PointField.FLOAT32, 1)
        ]
        
        msg.point_step = 16  # 4 * 4 bytes
        msg.row_step = msg.point_step * len(points)
        msg.height = 1
        msg.width = len(points)
        
        # 序列化点云数据
        msg.data = []
        for point in points:
            # 添加x, y, z坐标和强度
            msg.data.extend(struct.pack('ffff', point[0], point[1], point[2], intensity))
        
        # 发布消息
        publisher.publish(msg)
    
    def run(self):
        """运行节点"""
        rospy.loginfo("Grid Ground Filter node is running...")
        rospy.spin()

if __name__ == '__main__':
    try:
        filter_node = GridGroundFilter()
        filter_node.run()
    except rospy.ROSInterruptException:
        pass 